# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: dongkwan

import fire
import json
import os
import json
import pandas as pd

from collections import defaultdict


def run(
  tier=3,
  only_primary=True,
):
  if only_primary:
    primary_key = "primary_true"
  else:
    primary_key = "primary_false"

  export_path = os.path.expandvars(f"$HOME/workspace/coin/data/coin2/quant/lukka_tier{tier}_{primary_key}_category.json")
  df = pd.read_csv(os.path.expandvars("$HOME/workspace/coin/data/coin2/quant/lukka/lukka_sector_sample_raw.csv"))
  sector_map = defaultdict(set)

  for i in range(len(df)):
    row = df.iloc[i]
    if only_primary and not row["is Primary Sector"]:
      continue
    symbol = row["Presto Labs Provided Symbol"]
    sector = row[f"tier {tier} code"]

    sector_map[str(sector)].add(symbol)

  res = {
    "categories": sorted(list(sector_map.keys())),
    "infos": {}
  }
  for k in sorted(sector_map.keys()):
    res["infos"][k] = sorted(sector_map[k])
  with open(export_path, "w") as f:
    json.dump(res, f, indent=2)

  print(sector_map)
  print(f"Dump finished to {export_path}")


if __name__ == "__main__":
  fire.Fire(run)
